import msoffcrypto
import itertools
import string
from tqdm import tqdm
import os
import io

def brute_force_decrypt(file_path, password_length, charset):
    """暴力破解加密文档
    
    :param file_path: 文档路径
    :param password_length: 密码长度
    :param charset: 密码可能的字符集
    """
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"错误：文件不存在 - {file_path}")
        return False
    
    # 检查加密状态
    if not check_encryption(file_path):
        print("文件未加密或不受支持")
        return False
    
    # 计算总尝试次数
    total = len(charset) ** password_length
    print(f"字符集大小: {len(charset)}, 总尝试次数: {total:,}")
    
    passwords = itertools.product(charset, repeat=password_length)
    
    # 进度条显示
    success = False
    with tqdm(total=total, desc=f"破解 {password_length}位密码") as pbar:
        for pwd_tuple in passwords:
            password = ''.join(pwd_tuple)
            try:
                with open(file_path, "rb") as f:
                    office_file = msoffcrypto.OfficeFile(f)
                    office_file.load_key(password=password)
                    
                    # 尝试解密
                    decrypted_data = io.BytesIO()
                    office_file.decrypt(decrypted_data)
                    
                    # 验证解密内容
                    if decrypted_data.getbuffer().nbytes > 0:
                        with open("decrypted.docx", "wb") as out_file:
                            out_file.write(decrypted_data.getbuffer())
                        print(f"\n✅ 解密成功！密码是: [{password}]")
                        success = True
                        break
            except Exception as e:
                # 忽略错误继续尝试
                pass
            finally:
                pbar.update(1)
    
    if not success:
        print("\n❌ 未找到正确密码，请尝试其他字符集或长度")
    return success

def check_encryption(file_path):
    """检查文件是否加密"""
    try:
        with open(file_path, "rb") as f:
            office_file = msoffcrypto.OfficeFile(f)
            return office_file.is_encrypted()
    except Exception as e:
        print(f"加密检查错误: {str(e)}")
        return False

# ======================== 使用示例 ========================
if __name__ == "__main__":
    file_path = r"/Users/74670cqy/Desktop/001.doc"
    
    # 定义可能的字符集层级
    charsets = {
        "纯数字": string.digits,  # 10个字符
        "数字+小写字母": string.digits + string.ascii_lowercase,  # 36个字符
        "数字+大小写字母": string.digits + string.ascii_letters,  # 62个字符
        "常见符号": string.digits + string.ascii_letters + "!@#$%^&*",  # 70个字符
        "全字符": string.printable.strip()  # 100个字符
    }
    
    password_length = 4
    
    print("="*50)
    print(f"开始破解文档: {file_path}")
    print(f"密码长度: {password_length}位")
    print("="*50)
    
    # 按层级尝试不同字符集
    for name, charset in charsets.items():
        print(f"\n⚡ 尝试字符集: {name} ({len(charset)}个字符)")
        if brute_force_decrypt(file_path, password_length, charset=charset):
            break